{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_huggingface import HuggingFaceEmbeddings \n",
    "\n",
    "import os\n",
    "import chromadb\n",
    "from langchain_chroma import Chroma\n",
    "from langchain_community.document_loaders import PyPDFLoader\n",
    "from langchain.text_splitter import RecursiveCharacterTextSplitter\n",
    "\n",
    "from transformers import pipeline\n",
    "from langchain_huggingface.llms import HuggingFacePipeline\n",
    "from transformers import AutoModelForSeq2SeqLM, AutoTokenizer\n",
    "\n",
    "import pandas as pd\n",
    "from langchain_core.prompts import ChatPromptTemplate\n",
    "from langchain_core.prompts import PromptTemplate\n",
    "from langchain_core.prompts import FewShotPromptTemplate\n",
    "from langchain.chains import create_retrieval_chain\n",
    "from langchain.chains.combine_documents import create_stuff_documents_chain\n",
    "\n",
    "from typing import List,Any\n",
    "from langchain_core.vectorstores.base import VectorStore\n",
    "from langchain_core.documents import Document\n",
    "from nltk.tokenize import sent_tokenize\n",
    "import numpy as np\n",
    "\n",
    "from langchain_core.retrievers import BaseRetriever\n",
    "\n",
    "from dotenv import load_dotenv\n",
    "\n",
    "import json\n",
    "\n",
    "from collections import defaultdict\n",
    "import ast\n",
    "\n",
    "from langchain_core.output_parsers import JsonOutputParser\n",
    "from pydantic import BaseModel, Field\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "embedding_model_name = 'sentence-transformers/all-MiniLM-L6-v2'\n",
    "embedding_model = HuggingFaceEmbeddings(model_name=embedding_model_name,model_kwargs={'device':'cuda'}) \n",
    "max_length = embedding_model._client.tokenizer.model_max_length - 50\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "if not chatmodel change to Chatprompttempelate instead of PromptTempelate\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "checkpoint = \"\"\n",
    "device = \"cuda\" # for GPU usage or \"cpu\" for CPU usage\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(checkpoint)\n",
    "eval_model = AutoModelForSeq2SeqLM.from_pretrained(checkpoint,\n",
    "                                              torch_dtype=torch.float16,\n",
    "                                              trust_remote_code=True).to(device)\n",
    "\n",
    "pipe = pipeline(\n",
    "    \"text2text-generation\",\n",
    "    model=eval_model,\n",
    "    tokenizer=tokenizer,\n",
    ")\n",
    "\n",
    "llm = HuggingFacePipeline(\n",
    "    pipeline=pipe\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "db_path = os.path.join(os.getcwd(),\"chroma_db\")\n",
    "try:\n",
    "    chroma_client = chromadb.PersistentClient(path=db_path)\n",
    "    vector_store = chroma_client.delete_collection(name='Book_embeddings')\n",
    "except:\n",
    "    pass\n",
    "\n",
    "pdf_paths = ['books/12_lead_ecg_the_art_of_interpretation.pdf','books/jane-huff-ecg-workout-exercises-in-arrhythmia-interpretation.pdf']\n",
    "\n",
    "\n",
    "split_docs= []\n",
    "for pdf in pdf_paths:\n",
    "    loader = PyPDFLoader(os.path.join(os.getcwd(),pdf))\n",
    "    docs = loader.load()\n",
    "    text_splitter = RecursiveCharacterTextSplitter(chunk_size = max_length,chunk_overlap = max_length//5,length_function = len)\n",
    "    documents = text_splitter.split_documents(documents=docs)\n",
    "    split_docs.extend(documents)\n",
    "    \n",
    "# print(split_docs)import pandas as pd\n",
    "vector_store = Chroma.from_documents(split_docs,embedding=embedding_model,persist_directory=db_path,collection_name='Book_embeddings')\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "feature extraction and refine feature selection"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "retriever = vector_store.as_retriever()\n",
    "disease = \"Myocardial infraction\"\n",
    "query = f\"How to interpret {disease} related arrhythmias in a 12-lead ecg?(No Markdown)\"\n",
    "relevant_docs = retriever.invoke(query)\n",
    "docs = [doc.page_content for doc in relevant_docs]\n",
    "# print(len(docs))\n",
    "print(docs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "retriever = vector_store.as_retriever()\n",
    "diseases = {\"ST/T segment change (STTC)\", \"myocardial infarction(MI)\", \"conduction disturbance (CD)\", \"hypertrophy(HYP)\"}\n",
    "diagnosis_guidance = defaultdict(str)\n",
    "guiding_prompt = ChatPromptTemplate[(\n",
    "                (\"system\",\"You are an expert literature retriever\"),\n",
    "                (\"human\",\"Create a diagnosis guidance to interpret {disease} related arrhytmias in a 12-lead ecg system.\"),\n",
    "                )]\n",
    "create_diagnosis_guidance = create_stuff_documents_chain(llm=llm,prompt = guiding_prompt,)\n",
    "\n",
    "for disease in diseases:\n",
    "    query = f\"How to interpret {disease} related arrhythmias in a 12-lead ecg? \"\n",
    "    retrieved_context = retriever.invoke(query)\n",
    "    docs = [doc.page_content for doc in docs]\n",
    "    # See output format\n",
    "    diagnosis_guidance[disease] = create_diagnosis_guidance.invoke({\"context\":docs,\"disease\":disease})\n",
    "retrieved_context_NORM = retriever.invoke(\"How to interpret if a person does not have any heart arrhytmia using 12-lead ecg system\")\n",
    "NORM_docs = [doc.page_content for doc in retrieved_context_NORM]\n",
    "diagnosis_guidance[\"NORM\"] = create_diagnosis_guidance(llm = llm, prompt = \"Create a diagnosis guidance to interpret if a person does not have any related arrhytmias in a 12-lead ecg system \")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "subclass_to_superclass = {\n",
    "    \"NORM\": \"NORM\",\n",
    "    \"NSR\": \"NORM\",  # Normal Sinus Rhythm\n",
    "    \"NML\": \"NORM\",  # Normal ECG\n",
    "    \"LAFB/LPFB\": \"CD\", \"IRBBB\": \"CD\", \"ILBBB\": \"CD\", \"CLBBB\": \"CD\", \"CRBBB\": \"CD\", \"_AVB\": \"CD\", \"IVCB\": \"CD\", \"WPW\": \"CD\",\n",
    "    \"LVH\": \"HYP\", \"RHV\": \"HYP\", \"LAO/LAE\": \"HYP\", \"RAO/RAE\": \"HYP\", \"SEHYP\": \"HYP\",\n",
    "    \"AMI\": \"MI\", \"IMI\": \"MI\", \"LMI\": \"MI\", \"PMI\": \"MI\",\n",
    "    \"ISCA\": \"STTC\", \"ISCI\": \"STTC\", \"ISC_\": \"STTC\", \"STTC\": \"STTC\", \"NST_\": \"STTC\"\n",
    "}\n",
    "\n",
    "statements_path = os.path.join(os.getcwd(),\"ptb-xl+/labels/ptbxl_statements.csv\")\n",
    "\n",
    "statements_df = pd.read_csv(statements_path)\n",
    "\n",
    "count = {\n",
    "    \"NORM\": 0,\n",
    "    \"CD\": 0,\n",
    "    \"HYP\": 0,\n",
    "    \"MI\": 0,\n",
    "    \"STTC\": 0\n",
    "}\n",
    "\n",
    "diagnosis_statements = defaultdict(list)\n",
    "for idx, row in statements_df.iterrows():\n",
    "    ecg_id = row['ecg_id']\n",
    "    statements = ast.literal_eval(row['scp_codes'])\n",
    "    for statement in statements:\n",
    "        if statement[0] in subclass_to_superclass:\n",
    "            if(count[subclass_to_superclass[statement[0]]] < 10):\n",
    "                count[subclass_to_superclass[statement[0]]] += 1\n",
    "                diagnosis_statements[ecg_id].append(subclass_to_superclass[statement[0]])\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class diagnosisOutput(BaseModel):\n",
    "    disease_name : str\n",
    "    Result: bool\n",
    "    Explanation: str"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def diagnose(features:dict, disease:str):\n",
    "    feature_prompt = ChatPromptTemplate[(\n",
    "        (\"system\",\"You are an ECG expert and can calculate fiducial points and segements across different leads.\"),\n",
    "        (\"human\",\"Retrieve the relevant information from {features} required to interpret {disease} heart arrhythmias. Use this {diagnosis_guidance}\")\n",
    "    )]\n",
    "    patient_features = format_feature_flat(features)\n",
    "    create_feature_prompt = create_stuff_documents_chain(llm = llm, prompt = feature_prompt)\n",
    "    feature_prompt = create_feature_prompt.invoke({\n",
    "                    \"features\": patient_features,\n",
    "                    \"disease\": disease,\n",
    "                    \"diagnosis_guidance\": diagnosis_guidance[disease]\n",
    "                    \n",
    "                })\n",
    "    diagnosis_prompt = ChatPromptTemplate.from_template(\"\"\"\n",
    "        You are a medical expert in heart arrhythmia diagnosis.\n",
    "\n",
    "        DISEASE TO DIAGNOSE: {input}\n",
    "        MEDICAL CONTEXT: {context}\n",
    "        PATIENT'S MEASURED FEATURES: {patient_features}\n",
    "\n",
    "        Based on medical literature and the patient's measurements, determine whether the patient has {input}.\n",
    "\n",
    "        Please return ONLY a JSON object with the following fields:\n",
    "\n",
    "        - disease_name: Name of the disease (must match the input)\n",
    "        - Result: true or false, depending on whether the patient has the disease\n",
    "        - Explanation: Justification for the decision\n",
    "\n",
    "        Example output:\n",
    "        {{\n",
    "        \"disease_name\": \"NORM\",\n",
    "        \"Result\": False,\n",
    "        \"Explanation\": \"The ECG shows significant abnormalities including ST segment elevation, inverted T waves, and abnormal RS complexes, indicating the presence of arrhythmias and excluding a normal ECG\"\n",
    "        }}\n",
    "\n",
    "        ⚠️ Do not add any text before or after the JSON block.\n",
    "        Return only the JSON in valid format.\n",
    "        \"\"\")\n",
    "    parser = JsonOutputParser(pydantic_object=diagnosisOutput)\n",
    "    final_chain = diagnosis_prompt | llm | parser\n",
    "    \n",
    "    result = final_chain.invoke({\n",
    "        \"input\": disease,\n",
    "        \"context\": diagnosis_guidance[disease],\n",
    "        \"patient_features\": patient_features\n",
    "    })\n",
    "    return result\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "features_df = pd.read_csv(os.path.join(os.getcwd(),\"ptb-xl+/features/old/12sl_features.csv\"))\n",
    "tp = 0\n",
    "tn = 0\n",
    "fp = 0\n",
    "fn = 0\n",
    "df = pd.DataFrame()\n",
    "for ecg_id, statements in diagnosis_statements.items():\n",
    "    for _,features in features_df.iterrows():\n",
    "        if features['ecg_id'] == ecg_id:\n",
    "            results = []\n",
    "            for disease in diseases:\n",
    "                response = diagnose(features=features,disease=disease)\n",
    "                if(response.Result==True and disease in statements):\n",
    "                    tp +=1\n",
    "                elif(response.Result==True):\n",
    "                    tn +=1\n",
    "                elif(response.Result==False and disease in statements):\n",
    "                    fn +=1\n",
    "                else:\n",
    "                    fp +=1\n",
    "         "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "accuracy = (tp + tn) / (tp + tn + fp + fn) if (tp + tn + fp + fn) != 0 else 0\n",
    "precision = tp / (tp + fp) if (tp + fp) != 0 else 0\n",
    "recall = tp / (tp + fn) if (tp + fn) != 0 else 0\n",
    "f1_score = 2 * (precision * recall) / (precision + recall) if (precision + recall) != 0 else 0\n",
    "print(\"Accuracy : \", accuracy)\n",
    "print(\"Precision : \", precision)\n",
    "print(\"Recall : \", recall)\n",
    "print(\"F1 Score : \", f1_score)\n",
    "\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
